import KMeans
import ExperimentalAnalyzer as analyze
import sys
import os
from PIL import Image
from matplotlib import *
import pylab
from pylab import *
import Well


""" 
This is the main module. Run this from the command line for an interactive run. The ExperimentRunner
module is responsible for taking the user-supplied parameters and running the analysis then
dumping the data to the specified directory.

Warning: Research code. Very little (to no) error checking is done on the input. This will
break easily. Be careful.
"""
def interactive():
    """ This method deals with getting user I/O. (Driver for running the program interactively.) """
    print "Welcome to BacEsperanto ((BETA))"
    print "Developed at Tufts University\n"
    # Global reference to these for experimental runs.
    expPath, w, greenImgDir = None, None, None
    opt = getAndVerifyInput("Type 'N' to start a new experiment environment, or 'O' to open an existing one.", "Invalid selection. Try again.", lambda x: not x=='O' and not x=='o' and not x=='N' and not x=='n', "")
    if opt =='N' or opt=='n': 
        print "\nCreating a new experiment environment.\n"
        expDir = raw_input("Enter top-level directory for experiment output.\n>>")
        if os.path.isdir(expDir):
            print "Warning: directory already exists. Files might be overwritten."
        else:
            try:
                os.makedirs(expDir)
                print "Creating directory.."
            except:
                print "Couldn't make directory: " + expDir + "!"
                return
        print "Output directory for data set to: " + expDir + "\n"
        
        redIn = getAndVerifyInput("Enter the path to the Syto (red bacteria) image", "No image file found at specified path, please try again.",lambda x:not os.path.exists(x), "Red (syto) image set.\n") 
        sytosImg = Image.open(redIn)
           
        greenImgDir = getAndVerifyInput("Enter the path to the directory containing the Green bacteria images", "No directory found at specified path. Please enter a new path.", lambda x:not os.path.exists(x) or not os.path.isdir(x), "Green image directory set.")
        greensList = os.listdir(greenImgDir)
        greensList.sort()
        for g in greensList:
            print "Found green image: " + g
        print "T-Not image has been set to: " + greensList[0] + "\n"
        tNot = greensList[0]
        bgPath = getAndVerifyInput("Enter the path to the directory containing background images.", "No directory was found at the specified path. Please enter a new path.", lambda x:not os.path.isdir(x), "Background image directory set.\n")
        bgImages = os.listdir(bgPath)
        print "Background images found:"
        print bgImages
        print "Building background matrix..."
        picklePath = expDir + "\\bgMatrix.pickle"
        print "Pickling matrix to: " + picklePath
        m = analyze.buildBGMatrix([bgPath+"\\"+bgImage for bgImage in bgImages], writeOutMatrix=picklePath) 
        
        print "Background matrix built and pickled.\n"
        whPath = getAndVerifyInput("Enter path to well (white light) image.", "No image file found at the specified path. Please enter a new path to the white light image.", lambda x:not os.path.exists(x), "Well map image set.\n")
        wellMapImgOut = expDir+"\\well_map.tiff"
        wellMapPickleOut = expDir+"\\well_map.pickle"
        looksGood = False
        # initially we set the threshold to 2.2 (this is somewhat arbitrary) 
        # we bring it up and down according the difference in the discovered well
        # count and the expected well count.
        whThresh, w, adjustThreshold = 2,None,.2
        alreadyTried = []
        while not looksGood:
            # 2/6/07 -- automating the process of well map generation 
            # We might change this back to user-tuned parameterizing
            #whThresh = raw_input("Enter a threshold scalar for building the well map (typically ~2). You can change this afterwards if the map does not look correct.\n>>")
            wellCountShouldBe = 17500 # This is a constant from Ragnhild
            lax = 1000
            whThresh = whThresh+adjustThreshold
            if whThresh in alreadyTried:
                # If we've already tried a threshold that we've gotten back to,
                # it means we need to 'zoom in' and adjust by smaller values.
                # Note that we decrease exponentially.
                adjustThreshold = .5*adjustThreshold
                whThresh = whThresh+adjustThreshold
            alreadyTried.append(whThresh)
            print "Trying scalar threshold: " + str(whThresh)
            whProcessedOut = expDir+"\\WH_processed.tiff"
            #whThresh = float(whThresh[0:len(whThresh)]) # @TODO should be checking input here
            print "Pre-processing white image.."
            # 2/2/07 -- Should we be using the bg matrix? 
            # It seems to help with the dust on the plate in constructing
            # a more accurate map.
            analyze.colourPixelsM(whPath, m, (255,255,255), whThresh, whProcessedOut)
            #analyze.colourPixels(whPath, whThresh, (255,255,255), out=whProcessedOut)
            #colourPixels(imgPath, threshold, outColour, out="coloredPic.tif", bgColour=(0,0,0)):
            print "Done. Building Well Map..."
            im =Image.open(whProcessedOut)
            # The magic numbers here (2, 1750) are 'green threshold' and 'red threshold', respectively.
            # These are actually not used and this should be removed from the WellMap constructor.
            w = analyze.imgToWellMap(im, m, 2, 1750, wellMapImgOut, wellMapPickleOut)
            print "Number of wells identified: "
            print len(w.wellList)
            print "Done." 
            if abs(len(w.wellList) - wellCountShouldBe) < lax:
                print "Image accepted."
                print "Displaying the (in silico) Well Map image"
                wMapImg = Image.open(wellMapImgOut)
                wMapImg.show()
                looksGood=True
            else:
                print "Image rejected."
                print "Adjusting threshold scalar..."
                if len(w.wellList) < wellCountShouldBe:
                    # Here we found too few wells, so let's try and 
                    # take the scalar down
                    adjustThreshold = -1*adjustThreshold
                    print "Found too few wells: adjustment scalar is: " + str(adjustThreshold)
                else: 
                    # Here we found too many; so the threshold needs to be higher
                    adjustThreshold = abs(adjustThreshold)
                    # Notice that 'adjustment' is not a word
                    print "Found too many wells: adjustment scalar is: " + str(adjustThreshold)
                looksGood = False
            #ans = raw_input("Does this image look correct? Y/N\n>>")
            #if ans == "Y" or ans=="y":
            #    looksGood=True
            #else:
            #    looksGood=False
        w.setTNot(Image.open(greenImgDir+"\\"+tNot))
        w.setSytosImg(sytosImg) 
        print "Experiment directories set: " 
        print "\tred (syto) image path: " + redIn
        print "\tgreen image directory path: " + greenImgDir
        saveIt = getAndVerifyInput("Do you want to save this experiment environment? Y/N", "Not a valid option, type 'Y' or 'N' please.", lambda x: not x=="Y" and not x=="N", "")
        if saveIt == "Y":
            outF = open(expDir+"\\expEnv.txt", 'w')
            # Obviously, this should be saved to a better format (e.g., XML).
            # This is down and dirty (... and not flexible.)
            outF.write(expDir+"\n")
            outF.write(greenImgDir+"\n")
            outF.write(wellMapPickleOut+"\n")
            outF.write(redIn+"\n")
            outF.write(greenImgDir+"\\"+tNot+"\n")
            outF.close()
            print "Environment saved to:"+ expDir+"\\expEnv.txt"
    else:
        path = getAndVerifyInput("Enter the path to the saved experiment (.txt) file.", "File doesn't exist. Try again.", lambda x:not os.path.exists(x), "")
        print "Opening " + path
        inF = open(path, 'r')
        lines = inF.readlines()
        print "Unpickling well at " + lines[2][0:len(lines[2])-1] + " ..."
        w = analyze.unpickleFromFile(lines[2][0:len(lines[2])-1])
        print "Well Map unpickled."
        expDir = lines[0][0:len(lines[0])-1]
        print "Experiment directory set: " + expDir
        greenImgDir = lines[1][0:len(lines[1])-1]
        print "Green image directory set: " + greenImgDir
        redIn = lines[3][0:len(lines[3])-1]
        print "Red (syto) image set: " + redIn
        sytosImg = Image.open(redIn)
        tNot = lines[4][0:len(lines[4])-1]
        print "TNot (for greens) image set: " +  tNot
        w.setTNot(Image.open(tNot))
        w.setSytosImg(sytosImg) 
        print "Experiment environment constructed.\n"
    expN = getAndVerifyInput("Enter a name for this experiment.", "", lambda x:False, "")
    expPath =expDir+"\\"+expN
    print "Results will be saved to " + expPath
    numRegions = getAndVerifyInput("Enter number of spatial regions for this run.", "Not a number! Please enter a numerical number of regions.", lambda x:not isAnInt(x), "")
    # This is safe because we check for numerical casting in the above call.
    nRegions = int(numRegions)
    regionsSet = []
    spatialBoundsAndThresholds = []
    for region in range(nRegions):
        upperLeft=getAndVerifyInput("Enter upper left coordinate for region " + str(region) + " as follows: x,y (e.g., 0,0)", "", lambda x: False, "")
        upperLeftCoords = upperLeft.split(",")
        lowerRight=getAndVerifyInput("Enter lower right coordinate for region " + str(region) + " as follows: x,y (e.g., 20,30)", "", lambda x: False, "")
        lowerRightCoords = lowerRight.split(",")
        curBounds = []
        curBounds.append(map(lambda x: int(x), upperLeftCoords))
        curBounds.append(map(lambda x: int(x), lowerRightCoords))
        print "Bounds specified for this region:"
        print curBounds
        iThresh = getAndVerifyInput("Enter red (sender) threshold for this region.", "Not a number! Please enter a numerical threshold.", lambda x:not isAnInt(x), "")
        spatialBoundsAndThresholds.append((curBounds, int(iThresh)))
    print "Spatial regions and associated intensity thresholds for senders: "
    print spatialBoundsAndThresholds
    
    numTimeSlices = getAndVerifyInput("\nEnter number of time slices for this run.", "Not a number! Please enter a numerical number of time slices.", lambda x:not isAnInt(x), "")
    numTS = int(numTimeSlices)
    timeSliceList = []
    for timeSlice in range(numTS):
        startTime = getAndVerifyInput("Enter start time for time slice " + str(timeSlice) + ".", "Not a number! Please enter a numerical time slice.", lambda x:not isAnInt(x), "")
        endTime = getAndVerifyInput("Enter end time for time slice " + str(timeSlice) + ".", "Not a number! Please enter a numerical time slice.", lambda x:not isAnInt(x), "")
        timeSliceList.append((int(startTime), int(endTime)))
    print "Time Slices: "
    print timeSliceList
    
    numIntensityBuckets = getAndVerifyInput("\nEnter number of intensity buckets for this run.", "Not a number! Please enter a numerical number of intensity buckets.", lambda x:not isAnInt(x), "")
    nIB = int(numIntensityBuckets)
    intensityBucketList = []
    for ib in range(nIB):
        lowIntens = getAndVerifyInput("Enter the lower bound intensity for intensity bucket " + str(ib) + ".", "Not a number! Please enter a numerical lower intensity bound.", lambda x:not isAnInt(x), "")
        highIntens = getAndVerifyInput("Enter the higher bound intensity for intensity bucket " + str(ib) + ".", "Not a number! Please enter a numerical higher intensity bound.", lambda x:not isAnInt(x), "")
        intensityBucketList.append((int(lowIntens), int(highIntens)))
    print "Intensity Buckets: "
    print intensityBucketList    
    
    numN = getAndVerifyInput("Enter neighborhood range (e.g., 16)", "Not a number! Please enter a numerical neighborhood.", lambda x:not isAnInt(x), "Number of neighbors set.")
    numN = int(numN)
    print "Running experiment..."
    runExp = ExperimentRunner(intensityBucketList, timeSliceList, spatialBoundsAndThresholds, expPath, w, greenImgDir, n=numN)
    print "Finished."
    
def isAnInt(s):
    """ True if s can be cast as an int """
    try:
        int(s)
        return True
    except:
        return False

def getAndVerifyInput(initialQ, errorQ, rejectionF, successStr):
    """ Queries the user with the initialQ String parameter; if rejectionF returns False when
    fed the user's input then the user-input is returned. Else the user is queried again (and
    so on.) """
    returnToken = raw_input(initialQ+"\n>>")
    while rejectionF(returnToken):
        returnToken = raw_input(errorQ+"\n>>")
    print successStr
    return returnToken
    
class ExperimentRunner:
    def __init__(self, intensityBoundList, imageStartAndEndList, boundsAndThresholds, expDir, wellMap, greenImgDir, n=32, normalizeVectors=True):
        """ This method automates the process of running a batch of experiments and outputting
        the data. It will build a directory structure corresponding to the input parameters, e.g.,
        for each intensity 'bucket' in the intensityBoundList parameter there will be a subdirectory
        with the data for the different boundsAndThresholds tuples. The easiest way to understand this method
        is to look at an example call:
    
        bounds1 = ((0,0), (95, 508))
        bounds2 = ((94,0),(288,508))
        bounds3 = ((288,0),(665, 508))
        i1 = 2250
        i2 = 1700
        i3 = 1500
        spatialBoundsAndThresholds=[(bounds1, i1), (bounds2, i2), (bounds3, i3)]
        runExp([(200,400),(400,600),(600,800)], [(0,11),(11,21),(21,31),(31,41),(41,51)], spatialBoundsAndThresholds, "C:\\TB_Data\\9-27-06\\experiments14.n16.3", w, greenImgDir)
        
        Note that we build the spatialBoundsAndThresholds first for simplicity. This list is a 
        collection of tuples; the first entry gives a region of the image and the second gives an 
        intensity threshold for wells to be considered 'reds'. 
        """
        print intensityBoundList
        print imageStartAndEndList
        print boundsAndThresholds
        print expDir
        deltaRedCount, centroidDist = [],[]
        try:
            os.makedirs(expDir)
        except:
            print "Couldn't make directory: " + expDir + ", assuming it already exists.."
        for spatialBounds in boundsAndThresholds:
            # For each-region and corresponding 'red' threshold...
            sPath = expDir+"\\region_bounded_"+str(spatialBounds[0][0])+"-"+str(spatialBounds[0][1])
            os.mkdir(sPath)
            bounds = spatialBounds[0]
            redThresh = spatialBounds[1]
            print "Running experiments for image bounded by: " + str(bounds)
            print "Red threshold for this region is: " + str(redThresh)
            
            """ CHANGES 2 12 07 """ 
            #talkers = wellMap.findAllTalkerWells(redThresh, boundsBox=bounds)
            """ This is where you broke it !!! swap the above line back in too fix - bcw """
            # 2/12/07 -- trying to do stuff here
            # magicFunction is from a linear regression done on training data
            magicFunction = lambda x: 243.39710501160735 + x*14222.893622832175 
            # Need to test this function
            concentration = .82766666 # whatever RAGS gives you
            talkers = wellMap.findAllTalkerWellsWithF(magicFunction, concentration)
            """ END CHANGES """ 
            
            print "Talkers (reds) list built for run." 
            alreadyFoundGreens = {}
            # Keep seperate 'already found' lists for each intensity bound.
            # Note that this means a well can pass through two levels of intensity
            # during one time period (it will be counted twice; once for each
            # level).
            for intensBound in intensityBoundList:
                alreadyFoundGreens[str(intensBound)]=[]
            print "Dictionary built:"
            print alreadyFoundGreens
            # This is for graphing. We keep a dictionary where the keys are the time segments
            # and the values are the number of 
            timeSliceDict = {}
            for tPeriod in imageStartAndEndList:
                timeSliceDict[tPeriod]={}
                tPath = sPath+"\\" +str(tPeriod[0])+"_"+str(tPeriod[1])
                os.mkdir(tPath)
                for intensityBound in intensityBoundList:
                    print "Processing set for time period: " + str(tPeriod) + ", intensity bounds " + str(intensityBound) + "..."
                    iPath = tPath + "\\intensity_bounds_"+str(intensityBound[0]) + "_" + str(intensityBound[1])
                    os.mkdir(iPath)
                    print "Time segment: " + str(tPeriod)
                    print intensityBound
                    responders = wellMap.findAllGreensUsingCutOff(greenImgDir, tPeriod[1], imgStart=tPeriod[0],boundsBox=bounds,cutOffFactorBounds=intensityBound,alreadyResponderLs=alreadyFoundGreens[str(intensityBound)])
                    
                    wellMap.drawMapImage(iPath + "\\well_map.tiff")
                    alreadyFoundGreens[str(intensityBound)].extend(responders)
                    
                    print "Greens list built for time period: " + str(tPeriod) + ", intensity bounds " + str(intensityBound)
                    gTotal = [wellMap.countNumOfRedWells(wellMap.getKNearestNeighbors(x, n)) for x in responders]
                    if not len(gTotal) == 0:
                        total = reduce(lambda x,y: x+y, gTotal)
                        avgRedCount = float(total)/float(len(responders))
                        print "Average number of red neighbors for greens: " + str(avgRedCount)
                        timeSliceDict[tPeriod][intensityBound]=avgRedCount
                    else:
                        timeSliceDict[tPeriod][intensityBound]=0
                    print "Clustering..."
                    outXL = iPath + "\\clustering_data.xls"
                    avgGraphOutDir = iPath
                    greenIntens = analyze.getWellIntensitiesOverTime(greenImgDir, responders, wellMap, imgStart=tPeriod[0],imgCutOff=tPeriod[1])
                    # 12/26/06 -- adding option to normalize.
                    if normalizeVectors:
                        print "Normalizing response vectors..."
                        #print "Original vectors: "
                        #print greenIntens
                        # Get all the norms from each found responder well
                        normScalarLs = map(lambda x:pylab.norm(x), greenIntens)
                        for k in range(len(greenIntens)):
                            greenIntens[k]=map(lambda x:x/normScalarLs[k], greenIntens[k])
                        print "Done. Vectors normalized. "
                    print "Clustering..."
                    clustering = analyze.genClusterDataTable(responders, greenIntens, wellMap, outXL) 
                    print "Done. Writing out data.."
                    if not clustering[0] == None:
                        analyze.writeAvgGraphOutForAllClusters(clustering[0], avgGraphOutDir, wellMap, greenImgDir, clustering[1],iStart=tPeriod[0],iCutOff=tPeriod[1])
                    # Clear out the greens 
                    wellMap.greenWells = []
                    avgRedCounts = clustering[1]
                    if len(avgRedCounts)==2:
                        deltaRedCount.append(abs(avgRedCounts[0]-avgRedCounts[1]))
                        # 1/23/07 -- Switching the call here to call the new dist between clsuters method
                        #centroidDist.append(.getEuclidDistance(clustering[0][0].centroid.coords, clustering[0][1].centroid.coords))
                        centroidDist.append(analyze.getEuclidDistanceBetweenClusters(clustering[0]))
                    else:
                        print "Warning: insufficient data; not plotting average red neighbor count vs. centroid distance for this time period/intensity bucket."
                    print "Generated data for " + str(tPeriod) + ", intensity bounds " + str(intensityBound)
                    print "Data written to: " + iPath + "\n"
            
            #### Graphing the time vs. avg red neighbor count ####
            # @TODO: this should be a method.
            colors = ('b','g','r','c','m','y','k')
            print "Generating time vs. avg number of red neighbors (for all wells that have fallen into the  " + str(len(intensityBoundList)) + " intensity bound buckets for region " + str(spatialBounds) + "..."
            print timeSliceDict
            xlabel('Time')
            ylabel('Number of Reds Within ' + str(n) + ' Well Radius')
            xAxis = timeSliceDict.keys()
            xAxis.sort()
            print "X-Axis (intensity bound buckets): " + str(xAxis)
            axisLbls, lines, labels = [""], [], []
            for r in xAxis:
                axisLbls.append(str(r[0])+"-"+str(r[1]))
            axisLbls.append("")
            # The two and blank entries in axisLbls are for spacing (aesthetic) purposes
            ind = arange(len(axisLbls))
    
            lines,labels=[],[]
            for i in range(len(intensityBoundList)):
                intensLine, xCoords = [], []
                k = 1
                for p in xAxis:
                    intensLine.append(timeSliceDict[p][intensityBoundList[i]])
                    xCoords.append(k)
                    k+=1
                plotted = scatter(xCoords, intensLine, c=colors[i])
                lines.append(plotted)
                labels.append(str(intensityBoundList[i]))
            xticks(ind, axisLbls)
            legend(lines, labels, 'best', shadow=True)
            savefig(sPath+"\\graph_at_time_slices")                                         
            print "Finished graphing... Figure saved.\n"
            clf()
            if len(deltaRedCount) != 0:
                print "Plotting average cluster centroid difference vs. difference in sender neighbor count."
                xlabel('Euclidean Distance Between Centroids')
                ylabel('Difference Between Sender Neighbor Count Over Clusters')
                scatter(centroidDist, deltaRedCount)
                savefig(sPath+"\\CentroidDistanceVsSenderCount")
                print "Finished graphing... Figure saved.\n"
                clf()
            # Clear out the reds
            talkers = []
            wellMap.redWells = []
        print "Experiment finished. Dumped information to: " + expDir
       
print "((starting up))..."
interactive()

